package adkr

import (
    "ebbflow/pkg/crypto"
    "encoding/json"
    "log"
    "strconv"
    "sync"
    "time"
)

// PredicateFunc is a function type for the MVBA's external validation logic.
type PredicateFunc func(value []byte, localSet map[int]struct{}) bool

// DefaultMVBAPredicate returns the built-in predicate equivalent to Python's default,
// with optional strong-check gating via IsStrongMVBAEnabled().
func DefaultMVBAPredicate(n, f int) PredicateFunc {
    return func(value []byte, localSet map[int]struct{}) bool {
        if len(value) == 0 { return false }
        var arr []int
        if err := json.Unmarshal(value, &arr); err != nil { return false }
        seen := make(map[int]struct{}, len(arr))
        for _, id := range arr {
            if _, ok := localSet[id]; !ok { return false }
            if _, dup := seen[id]; dup { return false }
            if IsStrongMVBAEnabled() && !IsDealerValid(id) { return false }
            seen[id] = struct{}{}
        }
        return len(seen) >= n - f
    }
}

// TThresholdMVBA is the Go equivalent of the python function.
func TThresholdMVBA(
	sid string,
	pid, n, f int,
	c []int,
	g *crypto.G1,
	inputCh <-chan []byte,
	decideCh chan<- []byte,
	msgInCh <-chan Message,
	msgOutCh chan<- Message,
	coinBr func(payload []byte),
	localSetCh <-chan int,
    predicate PredicateFunc,
    coinPKs []crypto.G1,
    coinSK crypto.ZR,
    thresholdPK *crypto.G1,
) {
	log.Printf("MVBA-SID:%s PID:%d Starting...", sid, pid)

    // --- 1. Internal Channels and Dispatcher ---
	rbcInChs := make([]chan Message, n)
	for i := 0; i < n; i++ {
		rbcInChs[i] = make(chan Message, 100)
	}
    // coin channel to capture coin-related messages if needed later
    coinInCh := make(chan Message, 100)

    go func() {
		for msg := range msgInCh {
			if msg.Tag == ADKR_MVBA && msg.Round >= 0 && msg.Round < n {
				rbcInChs[msg.Round] <- msg
            } else if msg.Tag == ADKG_COIN {
                coinInCh <- msg
			}
		}
	}()

	// --- 2. Synchronously populate localSet BEFORE starting validation ---
	localSet := make(map[int]struct{})
	log.Printf("MVBA-SID:%s PID:%d waiting to receive local set...", sid, pid)
	for id := range localSetCh {
		localSet[id] = struct{}{}
	}
	log.Printf("MVBA-SID:%s PID:%d local set received, size: %d", sid, pid, len(localSet))

    // If未提供谓词，则使用默认谓词（非空校验 + 快速模式下直接通过）
    if predicate == nil {
        // 默认谓词：值应为 JSON 数组[]int；长度>=n-f；元素去重且均在 localSet 中
        predicate = func(value []byte, localSet map[int]struct{}) bool {
            if len(value) == 0 { return false }
            var arr []int
            if err := json.Unmarshal(value, &arr); err != nil { return false }
        // 去重与成员校验 + ACSS验证通过的 dealer 校验
            seen := make(map[int]struct{}, len(arr))
            for _, id := range arr {
            if _, ok := localSet[id]; !ok { return false }
            if _, dup := seen[id]; dup { return false }
            if IsStrongMVBAEnabled() && !IsDealerValid(id) { return false }
                seen[id] = struct{}{}
            }
            return len(seen) >= n - f
        }
    }

    // --- 3. PHASE 1: BROADCAST PROPOSALS (Start N RBCs) ---
	rbcInputChs := make([]chan []byte, n)
	rbcOutputChs := make([]chan []byte, n)

    for i := 0; i < n; i++ {
        rbcInputChs[i] = make(chan []byte, 1)
        rbcOutputChs[i] = make(chan []byte, 1)
        go StartRBCWithMode(
            sid+"RBC"+strconv.Itoa(i), pid, n, f, i,
            rbcInputChs[i], rbcOutputChs[i], rbcInChs[i], msgOutCh,
        )
    }

	myProposal := <-inputCh
	rbcInputChs[pid] <- myProposal
	log.Printf("MVBA-SID:%s PID:%d Broadcasted its proposal via RBC.", sid, pid)

	// --- 4. PHASE 2: WAIT FOR N-F VALID PROPOSALS ---
	deliveredRBCOutputs := make([][]byte, n)
	var deliveredMutex sync.Mutex
	deliveredCount := 0
	nfDeliveredCond := sync.NewCond(&sync.Mutex{})

	for i := 0; i < n; i++ {
		go func(leaderID int) {
			val := <-rbcOutputChs[leaderID]
			for !predicate(val, localSet) {
				time.Sleep(10 * time.Millisecond)
			}

			deliveredMutex.Lock()
			if deliveredRBCOutputs[leaderID] == nil {
				deliveredRBCOutputs[leaderID] = val
				deliveredCount++
				log.Printf("MVBA-SID:%s PID:%d Validated proposal from %d (%d/%d received).", sid, pid, leaderID, deliveredCount, n-f)
				if deliveredCount >= n-f {
					nfDeliveredCond.Broadcast()
				}
			}
			deliveredMutex.Unlock()
		}(i)
	}

	nfDeliveredCond.L.Lock()
	for deliveredCount < n-f {
		nfDeliveredCond.Wait()
	}
	nfDeliveredCond.L.Unlock()
    log.Printf("MVBA-SID:%s PID:%d Received N-F valid proposals. Proceeding to coin/ABA phase.", sid, pid)

    // --- 5. PHASE 3: 使用 coin 驱动决策（与 Python 随机化流程对齐，快速模式则走确定性最小 leader） ---
    validIdx := make([]int, 0, n)
    for i := 0; i < n; i++ {
        if deliveredRBCOutputs[i] != nil { validIdx = append(validIdx, i) }
    }
    if len(validIdx) == 0 {
        // Fallback: 使用本地或首个可用提案
        if deliveredRBCOutputs[pid] != nil { decideCh <- deliveredRBCOutputs[pid]; return }
        for i := 0; i < n; i++ { if deliveredRBCOutputs[i] != nil { decideCh <- deliveredRBCOutputs[i]; return } }
        time.Sleep(50 * time.Millisecond)
        decideCh <- nil
        return
    }

    // 快速模式：直接选择最小 leader，确保可复现
    if IsFastBench() {
        chosen := validIdx[0]
        for _, idx := range validIdx { if idx < chosen { chosen = idx } }
        log.Printf("MVBA-SID:%s PID:%d [FAST] choose smallest valid leader %d", sid, pid, chosen)
        decideCh <- deliveredRBCOutputs[chosen]
        return
    }

    // 非快速模式：使用简化 ABA 多轮 + coin 驱动收敛。
    // 选一个初始估计：最小有效 leader 的奇偶作为二值；也可将 leader 集映射成 bit。
    initBit := 0
    if len(validIdx) > 0 && (validIdx[0]%2 == 1) { initBit = 1 }

    coinSID := sid + "-coin"
    tcp := NewThresholdCoinProtocol(coinSID, pid, n, f, 0, c, g, 0, coinPKs, coinSK, thresholdPK, coinInCh, msgOutCh, log.Default())
    coinFn := func() (int, error) { return tcp.Run(0) }

    aba := NewABAProtocol(sid+"-aba", pid, n, f, msgInCh, msgOutCh)
    bit, _ := aba.Run(initBit, coinFn)

    // 将 bit 映射为具体 leader 选择：按 bit 选择 validIdx 的一半区间
    half := len(validIdx)/2
    if half == 0 { half = 1 }
    chosen := validIdx[0]
    if bit == 0 { chosen = validIdx[0] } else {
        if len(validIdx) > half { chosen = validIdx[half] } else { chosen = validIdx[len(validIdx)-1] }
    }
    log.Printf("MVBA-SID:%s PID:%d ABA-decide bit=%d choose leader %d", sid, pid, bit, chosen)
    decideCh <- deliveredRBCOutputs[chosen]
    return
}
